#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <numeric>
#include <vector>
#include<fstream>
#include<iostream>

#include "common/tensorrt.hpp"
#include "common/check.hpp"
#include "cuda_util.h"
#include "head.hpp"


namespace dsgn{
    class HeadImpl : public Head{
        public:
            virtual ~HeadImpl(){
                if (box_preds_) checkRuntime(cudaFree(box_preds_));
                if (cls_preds_) checkRuntime(cudaFree(cls_preds_));
                if (dir_cls_preds_) checkRuntime(cudaFree(dir_cls_preds_));
            }

            bool init(const std::string& model){
                engine_ = TensorRT::load(model);
                if (engine_ == nullptr) return false;

                box_preds_shape_ = engine_->static_dims(1);
                cls_preds_shape_ = engine_->static_dims(2);
                dir_cls_preds_shape_ = engine_->static_dims(3);


                checkRuntime(cudaMalloc(&box_preds_, sizeof(nvtype::half) * std::accumulate(box_preds_shape_.begin(), box_preds_shape_.end(), 1, std::multiplies<int>{})));
                checkRuntime(cudaMalloc(&cls_preds_, sizeof(nvtype::half) * std::accumulate(cls_preds_shape_.begin(), cls_preds_shape_.end(), 1, std::multiplies<int>{})));
                checkRuntime(cudaMalloc(&dir_cls_preds_, sizeof(nvtype::half) * std::accumulate(dir_cls_preds_shape_.begin(), dir_cls_preds_shape_.end(), 1, std::multiplies<int>{})));
                return true;
            }
            
            virtual void print() override {engine_->print("Detection Head");}

            virtual void forward(const nvtype::half* spatial_feature_2d, void* stream=nullptr){
                engine_->forward({spatial_feature_2d, box_preds_, cls_preds_, dir_cls_preds_},
                                 static_cast<cudaStream_t>(stream));
            }

            virtual nvtype::half* box_preds() override {return box_preds_;}
            virtual nvtype::half* cls_preds() override {return cls_preds_;}
            virtual nvtype::half* dir_cls_preds() override {return dir_cls_preds_;}
            virtual std::vector<int> box_preds_shape() override {return box_preds_shape_;}
            virtual std::vector<int> cls_preds_shape() override {return cls_preds_shape_;}
            virtual std::vector<int> dir_cls_preds_shape() override {return dir_cls_preds_shape_;}

        private:
            std::shared_ptr<TensorRT::Engine> engine_;
            nvtype::half* box_preds_ = nullptr;
            nvtype::half* cls_preds_ = nullptr;
            nvtype::half* dir_cls_preds_ = nullptr;
            std::vector<int> box_preds_shape_, cls_preds_shape_, dir_cls_preds_shape_;
    };

    std::shared_ptr<Head> createHead(const std::string& model){
        std::shared_ptr<HeadImpl> instance(new HeadImpl());
        if (!instance->init(model)){instance.reset();}
        return instance;
    };
}